#!/usr/bin/env python3
import argparse
import grpc
import os
import sys
import time
from time import sleep

# Import P4Runtime lib from parent utils dir
sys.path.append(
    os.path.join(os.path.dirname(os.path.abspath(__file__)),
                 'utils/'))
import p4runtime_lib.bmv2
from p4runtime_lib.error_utils import printGrpcError
from p4runtime_lib.switch import ShutdownAllSwitchConnections
import p4runtime_lib.helper

class P4Controller:
    def __init__(self, p4info_file_path, bmv2_file_path):
        self.p4info_helper = p4runtime_lib.helper.P4InfoHelper(p4info_file_path)
        self.bmv2_file_path = bmv2_file_path
        self.switches = {}

    def create_switch_connection(self, name, address, device_id, log_file, retries=5, delay=2):
        """创建交换机连接，支持重试机制"""
        for attempt in range(retries):
            try:
                print(f"Attempting to connect to {name} at {address} (attempt {attempt + 1}/{retries})...")
                switch = p4runtime_lib.bmv2.Bmv2SwitchConnection(
                    name=name,
                    address=address,
                    device_id=device_id,
                    proto_dump_file=log_file
                )
                # 测试连接
                switch.MasterArbitrationUpdate()
                print(f"Successfully connected to {name}")
                self.switches[name] = switch
                return switch
            except grpc.RpcError as e:
                print(f"Failed to connect to {name} (attempt {attempt + 1}/{retries}): {e}")
                if attempt < retries - 1:
                    print(f"Retrying in {delay} seconds...")
                    time.sleep(delay)
                else:
                    print(f"Giving up on {name} after {retries} attempts")
                    raise

    def setup_switch(self, switch):
        """配置单个交换机"""
        try:
            # 安装P4程序
            switch.SetForwardingPipelineConfig(
                p4info=self.p4info_helper.p4info,
                bmv2_json_file_path=self.bmv2_file_path
            )
            print(f"Successfully configured pipeline on {switch.name}")
        except grpc.RpcError as e:
            print(f"Failed to setup switch {switch.name}: {e}")
            raise

    def write_ipv4_lpm_rule(self, switch, dst_addr, prefix_len, dst_mac, port):
        """写入IPv4 LPM规则"""
        try:
            table_entry = self.p4info_helper.buildTableEntry(
                table_name="MyIngress.ipv4_lpm",
                match_fields={
                    "hdr.ipv4.dstAddr": [dst_addr, prefix_len]
                },
                action_name="MyIngress.ipv4_forward",
                action_params={
                    "dstAddr": dst_mac,
                    "port": port
                }
            )
            switch.WriteTableEntry(table_entry)
            print(f"Added IPv4 LPM rule to {switch.name}: {dst_addr}/{prefix_len} -> MAC:{dst_mac} Port:{port}")
        except grpc.RpcError as e:
            print(f"Failed to add IPv4 LPM rule to {switch.name}: {e}")
            raise

    def write_acl_ip_rule(self, switch, dst_addr, mask, priority):
        """写入ACL IP规则"""
        try:
            table_entry = self.p4info_helper.buildTableEntry(
                table_name="MyIngress.acl_ip_ternary",
                match_fields={
                    "hdr.ipv4.dstAddr": [dst_addr, mask]
                },
                action_name="MyIngress.drop",
                priority=priority
            )
            switch.WriteTableEntry(table_entry)
            print(f"Added ACL IP rule to {switch.name}: {dst_addr}/{mask} (priority:{priority})")
        except grpc.RpcError as e:
            print(f"Failed to add ACL IP rule to {switch.name}: {e}")
            raise

    def write_acl_udp_rule(self, switch, dst_port, mask, priority):
        """写入ACL UDP规则"""
        try:
            table_entry = self.p4info_helper.buildTableEntry(
                table_name="MyIngress.acl_udp_ternary",
                match_fields={
                    "hdr.udp.dstPort": [dst_port, mask]
                },
                action_name="MyIngress.drop",
                priority=priority
            )
            switch.WriteTableEntry(table_entry)
            print(f"Added ACL UDP rule to {switch.name}: Port {dst_port}/{mask} (priority:{priority})")
        except grpc.RpcError as e:
            print(f"Failed to add ACL UDP rule to {switch.name}: {e}")
            raise

    def configure_s1_rules(self):
        """配置交换机s1的规则"""
        s1 = self.switches['s1']

        # IPv4 LPM规则
        ipv4_rules = [
            ("10.0.1.1", 32, "08:00:00:00:01:11", 1),
            ("10.0.2.2", 32, "08:00:00:00:02:22", 2),
            ("10.0.3.3", 32, "08:00:00:00:03:00", 3),
            ("10.0.4.4", 32, "08:00:00:00:04:00", 4)
        ]

        for dst_addr, prefix_len, dst_mac, port in ipv4_rules:
            self.write_ipv4_lpm_rule(s1, dst_addr, prefix_len, dst_mac, port)

        # ACL规则
        self.write_acl_ip_rule(s1, "10.0.4.4", 4294967295, 1)
        self.write_acl_udp_rule(s1, 80, 65535, 1)

    def configure_s2_rules(self):
        """配置交换机s2的规则"""
        s2 = self.switches['s2']

        ipv4_rules = [
            ("10.0.1.1", 32, "08:00:00:00:03:00", 4),
            ("10.0.2.2", 32, "08:00:00:00:04:00", 3),
            ("10.0.3.3", 32, "08:00:00:00:03:33", 1),
            ("10.0.4.4", 32, "08:00:00:00:04:44", 2)
        ]

        for dst_addr, prefix_len, dst_mac, port in ipv4_rules:
            self.write_ipv4_lpm_rule(s2, dst_addr, prefix_len, dst_mac, port)

    def configure_s3_rules(self):
        """配置交换机s3的规则"""
        s3 = self.switches['s3']

        ipv4_rules = [
            ("10.0.1.1", 32, "08:00:00:00:01:00", 1),
            ("10.0.2.2", 32, "08:00:00:00:01:00", 1),
            ("10.0.3.3", 32, "08:00:00:00:02:00", 2),
            ("10.0.4.4", 32, "08:00:00:00:02:00", 2)
        ]

        for dst_addr, prefix_len, dst_mac, port in ipv4_rules:
            self.write_ipv4_lpm_rule(s3, dst_addr, prefix_len, dst_mac, port)

    def configure_s4_rules(self):
        """配置交换机s4的规则"""
        s4 = self.switches['s4']

        ipv4_rules = [
            ("10.0.1.1", 32, "08:00:00:00:01:00", 2),
            ("10.0.2.2", 32, "08:00:00:00:01:00", 2),
            ("10.0.3.3", 32, "08:00:00:00:02:00", 1),
            ("10.0.4.4", 32, "08:00:00:00:02:00", 1)
        ]

        for dst_addr, prefix_len, dst_mac, port in ipv4_rules:
            self.write_ipv4_lpm_rule(s4, dst_addr, prefix_len, dst_mac, port)

    def initialize_switches(self):
        """初始化所有交换机连接和配置"""
        # 创建交换机连接
        switches_config = [
            ('s1', '127.0.0.1:50051', 0, 'logs/s1-p4runtime-requests.txt'),
            ('s2', '127.0.0.1:50052', 1, 'logs/s2-p4runtime-requests.txt'),
            ('s3', '127.0.0.1:50053', 2, 'logs/s3-p4runtime-requests.txt'),
            ('s4', '127.0.0.1:50054', 3, 'logs/s4-p4runtime-requests.txt')
        ]

        # 确保日志目录存在
        os.makedirs('logs', exist_ok=True)

        successful_switches = []
        for name, address, device_id, log_file in switches_config:
            try:
                switch = self.create_switch_connection(name, address, device_id, log_file)
                successful_switches.append(switch)
            except Exception as e:
                print(f"Warning: Could not connect to {name}, skipping...")
                continue

        if not successful_switches:
            raise Exception("Failed to connect to any switches!")

        # 配置所有成功连接的交换机
        for switch in successful_switches:
            try:
                self.setup_switch(switch)
            except Exception as e:
                print(f"Warning: Failed to setup {switch.name}, removing from active switches...")
                self.switches.pop(switch.name, None)

        print(f"Successfully configured {len(self.switches)} switches")

    def run(self):
        """运行控制器"""
        try:
            print("Starting P4 controller...")
            self.initialize_switches()

            if not self.switches:
                print("No switches available for configuration. Exiting.")
                return

            # 配置各交换机的规则
            if 's1' in self.switches:
                print("Configuring switch s1 rules...")
                self.configure_s1_rules()
            else:
                print("Skipping s1 configuration (not connected)")

            if 's2' in self.switches:
                print("Configuring switch s2 rules...")
                self.configure_s2_rules()
            else:
                print("Skipping s2 configuration (not connected)")

            if 's3' in self.switches:
                print("Configuring switch s3 rules...")
                self.configure_s3_rules()
            else:
                print("Skipping s3 configuration (not connected)")

            if 's4' in self.switches:
                print("Configuring switch s4 rules...")
                self.configure_s4_rules()
            else:
                print("Skipping s4 configuration (not connected)")

            print("Configuration completed!")

            # 保持运行以便查看状态
            print("Controller is running. Press Ctrl+C to exit.")
            while True:
                sleep(10)  # 延长睡眠时间

        except KeyboardInterrupt:
            print("\nShutting down controller.")
        except Exception as e:
            print(f"Controller error: {e}")
        finally:
            ShutdownAllSwitchConnections()

def check_environment():
    """检查环境是否准备就绪"""
    print("Checking environment...")

    # 检查必要的目录和文件
    if not os.path.exists('./utils'):
        print("Error: utils directory not found!")
        return False

    # 检查是否在正确的工作目录
    current_dir = os.path.basename(os.getcwd())
    if 'exercises' not in current_dir and 'tutorials' not in current_dir:
        print("Warning: You may not be in the correct working directory")

    return True

def main(p4info_file_path, bmv2_file_path):
    """主函数"""
    if not check_environment():
        print("Environment check failed. Please make sure you are in the correct directory.")
        return

    # 检查P4info和BMv2 JSON文件是否存在
    if not os.path.exists(p4info_file_path):
        print(f"Error: P4Info file not found: {p4info_file_path}")
        print("Please run 'make' to compile the P4 program first.")
        return

    if not os.path.exists(bmv2_file_path):
        print(f"Error: BMv2 JSON file not found: {bmv2_file_path}")
        print("Please run 'make' to compile the P4 program first.")
        return

    print("Starting P4 controller with:")
    print(f"  P4Info: {p4info_file_path}")
    print(f"  BMv2 JSON: {bmv2_file_path}")

    controller = P4Controller(p4info_file_path, bmv2_file_path)
    controller.run()

if __name__ == '__main__':
    parser = argparse.ArgumentParser(description='P4Runtime Controller')
    parser.add_argument('--p4info', help='p4info proto in text format from p4c',
                        type=str, action="store", required=False,
                        default='./build/acl.p4.p4info.txt')
    parser.add_argument('--bmv2-json', help='BMv2 JSON file from p4c',
                        type=str, action="store", required=False,
                        default='./build/acl.json')
    args = parser.parse_args()

    main(args.p4info, args.bmv2_json)
